import requests
import time
import zipfile
import os
import subprocess
import json
from pathlib import Path
from enum import Enum

class MinerUMode(Enum):
    """MinerU 运行模式"""
    LOCAL = "local"
    CLOUD = "cloud"

# 配置管理
class MinerUConfig:
    def __init__(self):
        # 通用配置
        self.mode = MinerUMode.CLOUD  # 默认使用云端模式
        self.timeout = 30
        self.retry_times = 10
        self.verbose = False
        
        # 云端配置
        self.api_key = None
        self.base_url = "https://mineru.net/api/v4"
        
        # 本地配置
        self.local_command = "magic-pdf"  # 本地 MinerU 命令
        self.local_output_dir = "./output"  # 本地输出目录
        self.local_config_path = "magic-pdf-minimal.json"  # 本地配置文件路径
        
        self._load_config()
    
    def _load_config(self):
        """加载配置，优先级：环境变量 > .env 文件 > 默认值"""
        # 尝试加载 .env 文件
        env_path = Path(__file__).parent.parent / ".env"
        if env_path.exists():
            self._load_env_file(env_path)
        
        # 通用配置
        mode_str = os.getenv('MINERU_MODE', 'cloud').lower()
        self.mode = MinerUMode.LOCAL if mode_str == 'local' else MinerUMode.CLOUD
        self.timeout = int(os.getenv('MINERU_TIMEOUT', self.timeout))
        self.retry_times = int(os.getenv('MINERU_RETRY_TIMES', self.retry_times))
        self.verbose = os.getenv('MINERU_VERBOSE', 'false').lower() == 'true'
        
        # 云端配置
        self.api_key = os.getenv('MINERU_API_KEY', self.api_key)
        self.base_url = os.getenv('MINERU_BASE_URL', self.base_url)
        
        # 本地配置
        self.local_command = os.getenv('MINERU_LOCAL_COMMAND', self.local_command)
        self.local_output_dir = os.getenv('MINERU_LOCAL_OUTPUT_DIR', self.local_output_dir)
        self.local_config_path = os.getenv('MINERU_LOCAL_CONFIG_PATH', self.local_config_path)
        
        # 如果是云端模式且没有 API 密钥，使用默认值（向后兼容）
        if self.mode == MinerUMode.CLOUD and not self.api_key:
            self.api_key = 'eyJ0eXBlIjoiSldUIiwiYWxnIjoiSFM1MTIifQ.eyJqdGkiOiIxODkwNDM1MCIsInJvbCI6IlJPTEVfUkVHSVNURVIiLCJpc3MiOiJPcGVuWExhYiIsImlhdCI6MTc0ODY5NjQ3MSwiY2xpZW50SWQiOiJsa3pkeDU3bnZ5MjJqa3BxOXgydyIsInBob25lIjoiMTU4MDE0MzgzODYiLCJvcGVuSWQiOm51bGwsInV1aWQiOiIwYmNiYmU5NC0yMTliLTRiMjMtYjFiZi05ZWE5ZDZjZjQ0OTMiLCJlbWFpbCI6IiIsImV4cCI6MTc0OTkwNjA3MX0.CflSvU6KF09ZxYLrW9uzXAjK24VYaDh3WMEaDDGeTbbU17n4n3pSbc4zMaWpchEAwPiwuVGahb_ak7XjgK0dpw'
            print("⚠️  使用默认 API 密钥，建议在 .env 文件中配置 MINERU_API_KEY")
    
    def _load_env_file(self, env_path):
        """简单的 .env 文件解析器"""
        try:
            with open(env_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith('#') and '=' in line:
                        key, value = line.split('=', 1)
                        key = key.strip()
                        value = value.strip().strip('"').strip("'")
                        os.environ[key] = value
        except Exception as e:
            print(f"警告：无法读取 .env 文件: {e}")
    
    def validate(self):
        """验证配置是否有效"""
        if self.mode == MinerUMode.CLOUD:
            if not self.api_key:
                raise ValueError("云端模式下 MinerU API 密钥未配置，请设置 MINERU_API_KEY 环境变量或在 .env 文件中配置")
        elif self.mode == MinerUMode.LOCAL:
            # 检查 mineru 包和主要 API 是否可用
            try:
                import mineru
                from mineru.backend.pipeline.pipeline_analyze import doc_analyze
                print("✅ MinerU Python API 可用（mineru.backend.pipeline.pipeline_analyze.doc_analyze）")
                return True
            except ImportError as e:
                print(f"❌ MinerU Python API 不可用: {e}")
                raise ImportError("本地模式需要安装 mineru 包和相关依赖。请运行 'pip install mineru' 并确保环境正确。")
            except Exception as e:
                print(f"⚠️  验证过程中出现异常: {e}")
                raise
        return True

# 全局配置实例
config = MinerUConfig()

# 向后兼容
api_key = config.api_key

def process_pdf_local(pdf_path, output_dir=None):
    """使用本地 MinerU 处理 PDF 文件"""
    # 直接调用 MinerU Python API 进行本地解析
    try:
        from mineru.cli.common import read_fn, prepare_env
        from mineru.data.data_reader_writer import FileBasedDataWriter
        from mineru.utils.enum_class import MakeMode
        from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
        from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
        from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
    except ImportError as e:
        raise ImportError(f"MinerU 本地解析依赖未安装: {e}")

    pdf_path = Path(pdf_path)
    if not pdf_path.exists():
        raise ValueError(f"PDF 文件不存在: {pdf_path}")

    if output_dir is None:
        output_dir = config.local_output_dir

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    file_name = pdf_path.stem
    pdf_bytes = read_fn(str(pdf_path))
    lang = "ch"  # 可根据需要调整

    # pipeline 解析
    infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = pipeline_doc_analyze(
        [pdf_bytes], [lang], parse_method="auto", formula_enable=True, table_enable=True)
    model_list = infer_results[0]
    images_list = all_image_lists[0]
    pdf_doc = all_pdf_docs[0]
    _lang = lang_list[0]
    _ocr_enable = ocr_enabled_list[0]
    local_image_dir, local_md_dir = prepare_env(str(output_path), file_name, "auto")
    image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
    middle_json = pipeline_result_to_middle_json(model_list, images_list, pdf_doc, image_writer, _lang, _ocr_enable, True)
    pdf_info = middle_json["pdf_info"]
    image_dir = str(os.path.basename(local_image_dir))
    # 生成md文件
    md_content_str = pipeline_union_make(pdf_info, MakeMode.MM_MD, image_dir)
    md_writer.write_string(f"{file_name}.md", md_content_str)
    # 生成json文件
    md_writer.write_string(f"{file_name}_middle.json", json.dumps(middle_json, ensure_ascii=False, indent=4))

    if config.verbose:
        print(f"本地 MinerU 处理成功，输出目录: {local_md_dir}")

    return str(local_md_dir)

def get_task_id(file_name):
    """获取 MinerU 解析任务 ID（仅云端模式）"""
    if config.mode != MinerUMode.CLOUD:
        raise ValueError("get_task_id 仅适用于云端模式，本地模式请使用 process_pdf")
    
    config.validate()  # 验证配置
    
    url = f'{config.base_url}/extract/task'
    header = {
        'Content-Type': 'application/json',
        "Authorization": f"Bearer {config.api_key}"
    }
    pdf_url = 'https://vl-image.oss-cn-shanghai.aliyuncs.com/pdf/' + file_name
    data = {
        'url': pdf_url,
        'is_ocr': True,
        'enable_formula': False,
    }

    try:
        res = requests.post(url, headers=header, json=data, timeout=config.timeout)
        if config.verbose:
            print(f"请求状态码: {res.status_code}")
            print(f"响应内容: {res.json()}")
        
        if res.status_code == 401:
            raise ValueError("MinerU API 密钥无效或已过期，请检查 MINERU_API_KEY 配置")
        elif res.status_code != 200:
            raise ValueError(f"MinerU API 请求失败，状态码: {res.status_code}, 响应: {res.text}")
        
        response_data = res.json()
        if "data" not in response_data:
            raise ValueError(f"MinerU API 响应格式错误: {response_data}")
        
        task_id = response_data["data"]['task_id']
        if config.verbose:
            print(f"获取到任务 ID: {task_id}")
        return task_id
        
    except requests.exceptions.Timeout:
        raise ValueError(f"MinerU API 请求超时（{config.timeout}秒）")
    except requests.exceptions.RequestException as e:
        raise ValueError(f"MinerU API 请求异常: {e}")

def get_result(task_id):
    """获取 MinerU 解析结果并下载"""
    url = f'{config.base_url}/extract/task/{task_id}'
    header = {
        'Content-Type': 'application/json',
        "Authorization": f"Bearer {config.api_key}"
    }

    retry_count = 0
    while retry_count < config.retry_times:
        try:
            res = requests.get(url, headers=header, timeout=config.timeout)
            
            if res.status_code == 401:
                raise ValueError("MinerU API 密钥无效或已过期")
            elif res.status_code != 200:
                raise ValueError(f"获取任务状态失败，状态码: {res.status_code}")
            
            response_data = res.json()
            if "data" not in response_data:
                raise ValueError(f"响应格式错误: {response_data}")
            
            result = response_data["data"]
            if config.verbose:
                print(f"任务状态: {result}")
            
            state = result.get('state')
            err_msg = result.get('err_msg', '')
            
            # 如果任务还在进行中，等待后重试
            if state in ['pending', 'running']:
                print("任务未完成，等待5秒后重试...")
                time.sleep(5)
                continue
            
            # 如果有错误，输出错误信息
            if err_msg:
                raise ValueError(f"MinerU 任务出错: {err_msg}")
            
            # 如果任务完成，下载文件
            if state == 'done':
                full_zip_url = result.get('full_zip_url')
                if full_zip_url:
                    local_filename = f"{task_id}.zip"
                    print(f"开始下载: {full_zip_url}")
                    
                    r = requests.get(full_zip_url, stream=True, timeout=config.timeout)
                    with open(local_filename, 'wb') as f:
                        for chunk in r.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                    print(f"下载完成，已保存到: {local_filename}")
                    
                    # 下载完成后自动解压
                    unzip_file(local_filename)
                else:
                    raise ValueError("未找到 full_zip_url，无法下载")
                return
            
            # 其他未知状态
            raise ValueError(f"未知任务状态: {state}")
            
        except requests.exceptions.Timeout:
            retry_count += 1
            if retry_count < config.retry_times:
                print(f"请求超时，重试 {retry_count}/{config.retry_times}")
                time.sleep(2)
            else:
                raise ValueError(f"请求超时，已重试 {config.retry_times} 次")
        except requests.exceptions.RequestException as e:
            retry_count += 1
            if retry_count < config.retry_times:
                print(f"请求异常，重试 {retry_count}/{config.retry_times}: {e}")
                time.sleep(2)
            else:
                raise ValueError(f"请求异常: {e}")

# 解压zip文件的函数
def unzip_file(zip_path, extract_dir=None):
    """
    解压指定的zip文件到目标文件夹。
    :param zip_path: zip文件路径
    :param extract_dir: 解压目标文件夹，默认为zip同名目录
    """
    if extract_dir is None:
        extract_dir = zip_path.rstrip('.zip')
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_dir)
    print(f"已解压到: {extract_dir}")

def process_pdf(pdf_input, output_dir=None):
    """
    统一的 PDF 处理接口，根据配置自动选择本地或云端模式
    
    Args:
        pdf_input: PDF 文件路径（本地模式）或文件名（云端模式）
        output_dir: 输出目录（可选）
    
    Returns:
        str: 处理结果路径或任务ID
    """
    config.validate()
    
    if config.mode == MinerUMode.LOCAL:
        print(f"使用本地 MinerU 处理: {pdf_input}")
        return process_pdf_local(pdf_input, output_dir)
    else:
        print(f"使用云端 MinerU 处理: {pdf_input}")
        task_id = get_task_id(pdf_input)
        print(f'任务 ID: {task_id}')
        get_result(task_id)
        return task_id

def set_mode(mode):
    """
    设置 MinerU 运行模式
    
    Args:
        mode: "local" 或 "cloud"
    """
    global config
    if mode.lower() == "local":
        config.mode = MinerUMode.LOCAL
        print("已切换到本地模式")
    elif mode.lower() == "cloud":
        config.mode = MinerUMode.CLOUD
        print("已切换到云端模式")
    else:
        raise ValueError("模式必须是 'local' 或 'cloud'")

def get_current_mode():
    """获取当前运行模式"""
    return config.mode.value

if __name__ == "__main__":
    # 示例用法
    print(f"当前模式: {get_current_mode()}")
    
    if config.mode == MinerUMode.LOCAL:
        # 本地模式示例
        pdf_path = './pdf_reports/【国信证券】工业与汽车触底反弹，良率影响短期营收.pdf'  # 本地PDF文件路径
        try:
            result = process_pdf(pdf_path)
            print(f"处理完成，输出目录: {result}")
        except Exception as e:
            print(f"处理失败: {e}")
    else:
        # 云端模式示例
        file_name = '【财报】中芯国际：中芯国际2024年年度报告.pdf'
        try:
            result = process_pdf(file_name)
            print(f"处理完成，任务ID: {result}")
        except Exception as e:
            print(f"处理失败: {e}")
